#!/usr/bin/env python3
import os,re
import numpy as np
from datetime import datetime,timedelta
import pendulum
import pandas as pd
from matplotlib.path import Path
import xarray as xr
import proplot as plot
#plot.rc['backend'] = 'Qt4Agg'
print(plot.rc['backend'])
import cartopy.crs as ccrs
from cartopy.io.shapereader import Reader
from cartopy.mpl.patch import geos_to_path
from cartopy.feature import ShapelyFeature
from shapely.geometry import MultiPolygon
import cmaps
import argparse
import concurrent.futures
from xml.dom import minidom

def parse_time(string):
    match = re.match(r'(\d{4}\d{2}\d{2})(\d{2})?(\d{2})?', string)
    if match.group(1) and match.group(2) and match.group(3):
        return pendulum.from_format(string, 'YYYYMMDDHHmm')
    if match.group(1) and match.group(2):
        return pendulum.from_format(string, 'YYYYMMDDHH')
    elif match.group(1):
        return pendulum.from_format(string, 'YYYYMMDD')

def parsexml(filename):
    print(f'[Notice]: Reading {filename}')
    xmldoc = minidom.parse(filename)
    itemlist = xmldoc.getElementsByTagName('R')
    ele = []
    for s in itemlist:
        ele.append( [ datetime( int(s.attributes['Year'].value),int(s.attributes['Mon'].value), \
            int(s.attributes['Day'].value), int(s.attributes['Hour'].value), ), \
            str(s.attributes['Station_Id_C'].value), float(s.attributes['Lat'].value), float(s.attributes['Lon'].value), \
            float(s.attributes['PRE_1h'].value) ] )
    # print(type(s.attributes['Hour'].value))

    return pd.DataFrame(ele, columns = ['Time', 'Station_Id_C', 'Lat', 'Lon', 'PRE_1h'])




def plot_pcp(pcp, df, fig_dir, output_prefix, fig_format):
    df = df[ ( df['PRE_1h'] != 999999.0 ) & ( df['PRE_1h'] != 999998) ]
    df = df.sort_values(by='PRE_1h')
    print(df[ df['PRE_1h'] > 50 ])
    # load china shapefile
    shp_file = f'{os.path.dirname(os.path.realpath(__file__))}/shapefiles/china.shp'
    shape_records = Reader(shp_file).records()
    chn_geoms = []
    for country in shape_records:
        name = country.attributes['FCNAME'].rstrip('\x00')
        chn_geoms += [country.geometry]
        if name == '甘肃省':
            geoms  = [ country.geometry ]
            gs_geoms = MultiPolygon([country.geometry])
            path   = Path.make_compound_path(*geos_to_path(geoms))
    chn_geoms = MultiPolygon(chn_geoms)
    # chinese font
    cnfont = {'fontname':'fangsong'}
    f, axs = plot.subplots(ncols=2, nrows=1, figsize=(16,8),proj='pcarree' )
    axs.format(
        labels=True,latlines=10, lonlines=10,lonlim=(pcp.lon.values[0],pcp.lon.values[-1]),latlim=(pcp.lat.values[0],pcp.lat.values[-1]),
        suptitle='1小时累积降水(mm)  {}'.format(pd.to_datetime(pcp.time.values).strftime('%Y-%m-%dT%H')),**cnfont
    )
 
    levels = [0,0.1,0.5,1,2,3,4,5,6,7,8,9,10,15,20,25,50]
    im = axs[0].contourf(pcp.lon, pcp.lat, pcp * 1000, levels=levels, cmap=cmaps.precip3_16lev, extend='max')
    axs[0].colorbar(im, loc='b', length=0.9)
    axs[0].format(title='分析', **cnfont)

    im = axs[1].scatter(df['Lon'].values, df['Lat'].values, marker='o', c=df['PRE_1h'].values, s=7, lw=0.2, cmap=cmaps.precip3_16lev, levels=levels,edgecolors='k')
    axs[1].colorbar(im, loc='b', length=0.9)
    axs[1].format(title='观测', **cnfont)
    shape_feature = ShapelyFeature(chn_geoms, ccrs.PlateCarree(), facecolor='none',edgecolor='k')
    axs.add_feature(shape_feature)

    if not os.path.exists(fig_dir):
        os.makedirs(fig_dir)
    _fig_path = os.path.join(fig_dir,'1hour_apcp_{}'.format(pd.to_datetime(pcp.time.values).strftime('%Y%m%d%H%M')))
    #plot.show()
    f.savefig(_fig_path + '.' + fig_format, dpi=150, bbox_inches='tight')
    return



if __name__ == '__main__':
    '''
    plot moc laps products ...
    '''
    parser = argparse.ArgumentParser(description='plot moc laps surface data.')
    parser.add_argument('-o', '--root-dir', dest='root_dir', default='/data/cma_moc/gfs-3km-prod', help='Root directory to store data.')
    parser.add_argument('--fig-dir', dest='fig_dir', default='/data/cma_moc/rtoas-figure/gfs-3km-prod', help='Root directory to store figure.')
    parser.add_argument('--output-prefix', dest='output_prefix', default='MOC_3KM_PCP', help='filename prefix.')
    parser.add_argument('-t', '--time', help='file time (YYYYMMDDHH[MM]).', type=parse_time)
    parser.add_argument('-f','--format', dest='format', default='png', help='figure format.')
    args = parser.parse_args()

    file_path = os.path.join(args.root_dir,args.time.format('YYYY'),args.time.format('YYYYMMDD'),f'{args.output_prefix}_{args.time.format("YYYYMMDDHHmm")}.nc')
    obs_path = os.path.join('/nas02/data/raw/cimiss_archive/SURF_CHN_MUL_HOR/',args.time.format('YYYYMMDD'),f'{args.time.format("YYYYMMDDHHmm")}.xml')
    # /nas02/data/raw/cimiss_archive/SURF_CHN_MUL_HOR/20201130/202011300300.xml
    df = parsexml(obs_path)
    try:
        ds = xr.open_dataset(file_path)
        ds.load()
    except:
        print(f'[Error]: not found {file_path}')
        exit(1)
    ds = ds.isel(time=0)
    plot_pcp(ds.r01, df, os.path.join(args.fig_dir,args.time.format('YYYY'),args.time.format('YYYYMMDD')), args.output_prefix, args.format)
